from mpmath import mp
import numpy as np

# 本方案与低尾数位 float 乘法的区别在于，在复杂度 O(n) > n ^ 2 的环节:
#   1. 线性重分布至 [1.0, 2.0 ** 0.5)，相乘后上限小于 2.0，保证了不会出现指数位进位；
#   2. 线性重分布至 [1.0, 2.0 ** 0.5)，都是正数，不必再考虑正负符号，usign ；
#   3. 用常数和指数位输出相乘改造而来的，有限次数加法，代替两个原始入参变量的尾数位相乘；
# 从而腾出了宝贵的数位，省去了对应的门电路
def g_gemm(i_arr_1, i_arr_2, i_dic_cnt = 512):
  # 设置精度
  mp.prec = i_dic_cnt * 2
  
  # 设置矩阵乘法的高、厚、宽
  h = i_arr_1.shape[-2]
  t = i_arr_1.shape[-1]  # i_arr_2.shape[-2]
  w = i_arr_2.shape[-1]

  # 入参
  # arr_1_ex = mp.matrix([[mp.rand() for _ in range(t)] for _ in range(h)])
  # arr_2_ex = mp.matrix([[mp.rand() for _ in range(w)] for _ in range(t)])
  arr_1_ex = mp.matrix(i_arr_1)  # h * t
  arr_2_ex = mp.matrix(i_arr_2)  # t * w
  
  # 设置求对数的底数和结果元素数值范围
  u = mp.mpf(2.0 ** (1.0 / i_dic_cnt))
  v = mp.mpf(2.0)
  
  # 准备线性重分布，求得入参的最大、最小、跨度
  arr_1_ex_max = np.array(mp.matrix(np.max(i_arr_1, axis = -1, keepdims = True)).tolist(), dtype=float)   # h * t   h * 1
  arr_2_ex_max = np.array(mp.matrix(np.max(i_arr_2, axis = -2, keepdims = True)).tolist(), dtype=float)   # t * w   1 * w
  arr_1_ex_min = np.array(mp.matrix(np.min(i_arr_1, axis = -1, keepdims = True)).tolist(), dtype=float)   # h * t   h * 1
  arr_2_ex_min = np.array(mp.matrix(np.min(i_arr_2, axis = -2, keepdims = True)).tolist(), dtype=float)   # t * w   1 * w
  arr_1_ex_ptp = arr_1_ex_max - arr_1_ex_min    # h * t   h * 1
  arr_2_ex_ptp = arr_2_ex_max - arr_2_ex_min    # t * w   1 * w
    
  m_ln2 = mp.log(u)
  
  # 求线性重分布的参数
  k1 = ((v ** 0.5) - 1.0) / arr_1_ex_ptp            # h * 1
  k2 = (((v ** 0.5) / u) - 1.0) / arr_2_ex_ptp      # 1 * w
  b1 = 1.0 - (k1 * arr_1_ex_min)                    # h * 1
  b2 = 1.0 - (k2 * arr_2_ex_min)                    # 1 * w
  
  # 线性重分布
  # uint 类型转换放在这一步也行，放在下一步也行。放在下一步的话，定点小数相加再四舍五入，能精确一些。
  arr_1_ex_kb = mp.matrix([[k1[i, 0] * arr_1_ex[i, j] + b1[i, 0] for j in range(t)] for i in range(h)])          # h * 1   h * t
  arr_2_ex_kb = mp.matrix([[k2[0, j] * arr_2_ex[i, j] + b2[0, j] for j in range(w)] for i in range(t)])          # 1 * w   t * w
  
  # 求对数
  ## arr_1_ex_log = mp.matrix(arr_1_ex_kb.rows, arr_1_ex_kb.cols, lambda i, j: mp.log(arr_1_ex_kb[i, j]) / m_ln2)
  ## arr_2_ex_log = mp.matrix(arr_2_ex_kb.rows, arr_2_ex_kb.cols, lambda i, j: mp.log(arr_2_ex_kb[i, j]) / m_ln2)
  arr_1_ex_log = np.uint16(np.round(np.array(mp.matrix([[mp.log(arr_1_ex_kb[i, j]) / m_ln2 for j in range(t)] for i in range(h)]).tolist(), dtype = float)))      # h * t
  arr_2_ex_log = np.uint16(np.round(np.array(mp.matrix([[mp.log(arr_2_ex_kb[i, j]) / m_ln2 for j in range(w)] for i in range(t)]).tolist(), dtype = float)))      # t * w
  
  # 用对数做加法指令，替换矩阵乘法的乘法指令
  # uint 类型 也可采用定点小数类型，小数点后一位小数
  # 朴素算法
  # tmp_a = np.sum(u ** (np.uint16(np.array(arr_1_ex_log)[:, :, np.newaxis] + np.array(arr_2_ex_log)[np.newaxis, :, :])), axis = -2, keepdims = False)   # h * w
  
  # ** 幂运算用下列方法分段替代(如下是分5段的例子): 8 位二进制常数 * 16 位 int，可以改造为少于8次的加法
  # select a_no,
  #   2.0 ^ (a_no :: float / a_i_dic_cnt), 1 + (a_no :: float / a_i_dic_cnt)
  # , ((2.0 ^ (a_no :: float / a_i_dic_cnt)) - (1 + (a_no :: float / a_i_dic_cnt))) * 8.0 -- 用位移运算和加法运算逼近
  # , 1.0 + (a_no :: float / a_i_dic_cnt) 
  #   + 
  #   case 
  #     when a_no between 0 and (5 * a_i_dic_cnt / 16) - 1                                          -- 0 and 80 - 1
  #       -- 两点定斜率位置(0 * a_i_dic_cnt / 16, 4 * a_i_dic_cnt / 16)(0, 64)(-0.0, -0.4863430799782318)
  #       then (((-0.4863430799782318 + 0.0) / (4 * a_i_dic_cnt / 16)) * (a_no - (0 * a_i_dic_cnt / 16)) -0.0) / 8.0  
  #     when a_no between (5 * a_i_dic_cnt / 16) and (a_i_dic_cnt / 2) - 1                            -- 80 and 128 - 1
  #       -- 两点定斜率位置(11 * a_i_dic_cnt / 32, 15 * a_i_dic_cnt / 32)(88, 120)(-0.5975923424661342, -0.6787209442893438)
  #       then (((-0.6787209442893438 + 0.5975923424661342) / (4 * a_i_dic_cnt / 32)) * (a_no - (11 * a_i_dic_cnt / 32)) - 0.5975923424661342) / 8.0  
  #     when a_no between (a_i_dic_cnt / 2) and (19 * a_i_dic_cnt / 32) - 1                           -- 128 and 152 - 1
  #       -- 两点定斜率位置(17 * a_i_dic_cnt / 32, 18 * a_i_dic_cnt / 32)(136, 144)(-0.6885535441836268, -0.6885650469709947, -0.6853908324840052)
  #       then (((-0.6853908324840052 + 0.6885535441836268) / (1 * a_i_dic_cnt / 32)) * (a_no - (17 * a_i_dic_cnt / 32)) - 0.6885535441836268) / 8.0
  #     when a_no between (19 * a_i_dic_cnt / 32) and (3 * a_i_dic_cnt / 4) - 1                       -- 152 and 192 - 1
  #       -- 两点定斜率位置(20 * a_i_dic_cnt / 32, 23 * a_i_dic_cnt / 32)(160, 184)(-0.662313396736474, -0.5839561747682804)
  #       then (((-0.5839561747682804 + 0.662313396736474) / (3 * a_i_dic_cnt / 32)) * (a_no - (20 * a_i_dic_cnt / 32)) - 0.662313396736474) / 8.0
  #     when a_no between (3 * a_i_dic_cnt / 4) and a_i_dic_cnt - 1                                   -- 192 and 256 - 1
  #       -- 两点定斜率位置(25 * a_i_dic_cnt / 32, a_i_dic_cnt)(200, 256)(-0.44998271701360437, -0.0)
  #       then (((0.0 + 0.44998271701360437) / (6 * a_i_dic_cnt / 32)) * (a_no - (26 * a_i_dic_cnt / 32)) - 0.44998271701360437) / 8.0
  #   end
  # from (select 65536.0 as a_i_dic_cnt) tb_a_i_dic_cnt(a_i_dic_cnt)
  # , generate_series(0, a_i_dic_cnt) tb_a(a_no)
  # order by a_no desc
  
  # 注意 np.round() 是银行家舍入法，四舍六入五成双，可能对准确度有影响  
  tmp_ae = np.uint16(np.round(np.array(arr_1_ex_log)[:, :, np.newaxis] + np.array(arr_2_ex_log)[np.newaxis, :, :]))
  tmp_a = np.zeros((h, t, w))
  for i in range(t) : 
    for j in range(h) : 
      for k in range(w) : 
        if 0 <= tmp_ae[j, i, k] < (5 * i_dic_cnt / 16) : 
          tmp_a[j, i, k] = 1.0 + (tmp_ae[j, i, k] / i_dic_cnt) + ((((-0.4863430799782318 - 0.0) / (4 * i_dic_cnt / 16)) * (tmp_ae[j, i, k] - (0 * i_dic_cnt / 16)) - 0.0) / 8.0)
        elif (5 * i_dic_cnt / 16) <= tmp_ae[j, i, k] < (i_dic_cnt / 2) :
          tmp_a[j, i, k] = 1.0 + (tmp_ae[j, i, k] / i_dic_cnt) + ((((-0.6787209442893438 + 0.5975923424661342) / (4 * i_dic_cnt / 32)) * (tmp_ae[j, i, k] - (11 * i_dic_cnt / 32)) - 0.5975923424661342) / 8.0)
        elif (i_dic_cnt / 2) <= tmp_ae[j, i, k] < (19 * i_dic_cnt / 32) :
          tmp_a[j, i, k] = 1.0 + (tmp_ae[j, i, k] / i_dic_cnt) + ((((-0.6853908324840052 + 0.6885535441836268) / (1 * i_dic_cnt / 32)) * (tmp_ae[j, i, k] - (17 * i_dic_cnt / 32)) - 0.6885535441836268) / 8.0)
        elif (19 * i_dic_cnt / 32) <= tmp_ae[j, i, k] < (3 * i_dic_cnt / 4) :
          tmp_a[j, i, k] = 1.0 + (tmp_ae[j, i, k] / i_dic_cnt) + ((((-0.5839561747682804 + 0.662313396736474) / (3 * i_dic_cnt / 32)) * (tmp_ae[j, i, k] - (20 * i_dic_cnt / 32)) - 0.662313396736474) / 8.0)
        elif (i_dic_cnt / 2) + (i_dic_cnt / 4) <= tmp_ae[j, i, k] < i_dic_cnt :
          tmp_a[j, i, k] = 1.0 + (tmp_ae[j, i, k] / i_dic_cnt) + ((((0.0 + 0.44998271701360437) / (6 * i_dic_cnt / 32)) * (tmp_ae[j, i, k] - (26 * i_dic_cnt / 32)) - 0.44998271701360437) / 8.0)
  
  tmp_a = np.sum(tmp_a, axis = -2, keepdims = False)

  # 还原公式验证如下: 
  #   with 
  #   cte_htw as 
  #   (
  #     select 
  #       3 as a_h
  #     , 4 as a_t
  #     , 5 as a_w
  #   ),
  #   cte_arr as 
  #   (
  #     select  
  #       sm_sc.fv_new_randn(0.0, 1.0, array[a_h, a_t]) as i_arr_1
  #     , sm_sc.fv_new_randn(0.0, 1.0, array[a_t, a_w]) as i_arr_2
  #     from cte_htw
  #   ),
  #   cte_kb as 
  #   (
  #     select 
  #       sm_sc.fv_new_randn(0.0, 1.0, array[a_h, 1]) as i_arr_1_k
  #     , sm_sc.fv_new_randn(0.0, 1.0, array[1, a_w]) as i_arr_2_k
  #     , sm_sc.fv_new_randn(0.0, 1.0, array[a_h, 1]) as i_arr_1_b
  #     , sm_sc.fv_new_randn(0.0, 1.0, array[1, a_w]) as i_arr_2_b
  #     from cte_arr, cte_htw
  #   ),
  #   cte_arr_kb as 
  #   (
  #     select 
  #       i_arr_1_k *` i_arr_1 +` i_arr_1_b as a_arr_1_kb
  #     , i_arr_2_k *` i_arr_2 +` i_arr_2_b as a_arr_2_kb
  #     from cte_arr, cte_kb
  #   ),
  #   cte_kb_back as 
  #   (
  #     select 
  #   	(a_arr_1_kb |**| a_arr_2_kb) 
  #   	-` 
  #   	-- (a_arr_1_kb |**| sm_sc.fv_repeat_axis_py(i_arr_2_b, 1, a_t)) 等价如下
  #   	((a_arr_1_kb |@+| array[1, a_t]) *` i_arr_2_b)
  #   	-`
  #   	-- (sm_sc.fv_repeat_axis_py(i_arr_1_b, 2, a_t) |**| a_arr_2_kb) 等价如下
  #   	(i_arr_1_b *` (a_arr_2_kb |@+| array[a_t, 1]))
  #   	+`
  #   	-- (sm_sc.fv_repeat_axis_py(i_arr_1_b, 2, a_t) |**| sm_sc.fv_repeat_axis_py(i_arr_2_b, 1, a_t)) 等价如下
  #   	((i_arr_1_b |**| i_arr_2_b) *` a_t)
  #   	as gemm_b
  #     from cte_kb, cte_arr_kb, cte_htw
  #   )
  #   select 
  #     (gemm_b /` i_arr_1_k /` i_arr_2_k) :: decimal[] ~=` 3
  #   , (i_arr_1 |**| i_arr_2) :: decimal[] ~=` 3 as gemm_ora
  #   from cte_arr, cte_kb, cte_kb_back

  # 准备线性重分布回去的其他多项式项
  tmp_b = np.sum(np.array(np.array(arr_1_ex_kb.tolist(), dtype=float)), axis = -1, keepdims = True) * b2     # h * w
  tmp_c = b1 * np.sum(np.array(np.array(arr_2_ex_kb.tolist(), dtype=float)), axis = -2, keepdims = True)     # h * w
  tmp_d = np.matmul(b1, b2) * t                                                         # h * w
  
  # 线性重分布还原
  tmp_e = tmp_a - (tmp_b + (tmp_c - tmp_d))
  return np.float16((tmp_e / k1) / k2)  # h * w

arr_1 = np.float16(np.random.randn(3, 5))
arr_2 = np.float16(np.random.randn(5, 2))
tmp_out = g_gemm(arr_1, arr_2)

# 与精确算法对比
print("近似算法:")
print(np.float16(tmp_out))
print("精确算法:")
print(np.float16(np.matmul(arr_1, arr_2)))

quit()